from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split


class DataPreparer:
    """
        获取digits数据集,
        @params
            train_size_:训练集比例
        @usage:
            data_preparaer = DataPreparer(train_size_=0.2)
            x_train,x_test,y_train,y_test = data_preparaer.get()
        @notice:
            返回的每个数据是1*64 而非 8*8
    """

    def __init__(self, train_size_) -> None:
        self.getData(train_size_)

    def getData(self, train_size_):

        # 加载数据
        digits_data, digits_target = load_digits(
            return_X_y=True)

        # 划分训练集和测试集
        self.x_train, self.x_test, self.y_train, self.y_test = \
            train_test_split(digits_data, digits_target,
                             test_size=1-train_size_, train_size=train_size_,
                             random_state=4, shuffle=True)

    # 外部获取训练集和测试集接口
    def get(self):
        return self.x_train, self.x_test, self.y_train, self.y_test
